import pickle
import csv

import numpy as np
import tensorflow as tf

if __name__ == "__main__":
    compare_result_path = "acos.csv"
    input_dict_list = [
        {"x": np.array([1, 2, 3], dtype=np.int32)},
        {"x": np.array([1, 2, 3], dtype=np.float16)},
        {"x": np.array([1, 2, 3], dtype=np.float32)},
    ]

    source_results = []
    follow_results = []

    for input_dict in input_dict_list:
        try:
            a = tf.math.asin(**input_dict)
        except Exception as e:
            print("[Run TF Model] Exception occurred~:\n", e)
            a = tf.constant([])
        print("Source Output:")
        print(a)
        # Save source output
        source_results.append(a.numpy())

        ragged_dict = dict(input_dict)
        for key in ragged_dict.keys():
            value = ragged_dict[key]
            if type(value) is np.ndarray:
                # Convert ndarray to ragged tensor
                if value.ndim == 1:
                    value = [value]
                ragged_dict[key] = tf.RaggedTensor.from_tensor(tf.convert_to_tensor(value))
        try:
            b = tf.math.asin(**ragged_dict)
        except Exception as e:
            print("[Run TF Model] Exception occurred~:\n", e)
            b = tf.constant([])
        print("Follow-up Output:")
        print(b)
        # Save follow-up output
        follow_results.append(b.numpy())

    # Compare source results and follow-up results
    compare_results = []
    for i in range(len(source_results)):
        a = source_results[i]
        b = follow_results[i]
        if np.allclose(a, b, atol=1.e-2, rtol=1.e-5, equal_nan=True):
            compare_results.append([i, "Consistent"])
        else:
            compare_results.append([i, "Inconsistent"])
    with open(compare_result_path, "w+", newline="") as f:
        w = csv.writer(f)
        w.writerows(compare_results)



